import json

from sqlalchemy.orm import Session
from config.database import get_db
from fastapi import WebSocket, WebSocketDisconnect, Depends
from websocket_module.service import ConnectionManager
from websocket_module import websocket_router
from config.logger import logging
from config.config import SECRET_KEY

logger = logging.getLogger(__name__)
manager = ConnectionManager()

@websocket_router.websocket("/{client_id}")
async def websocket_endpoint(
    websocket: WebSocket,
    client_id: str,
    db: Session = Depends(get_db)
):
    await websocket.accept()
    # 校验 client_id 是否合法（可选）
    if not client_id:
        await websocket.close(code=4000)
        return

    try:
        # 等待客户端发送认证信息
        auth_msg = await websocket.receive_json()
        if auth_msg.get("type") != "auth" or not auth_msg.get("token"):
            await websocket.send_json({"error": "Invalid auth message"})
            await websocket.close(code=4000)
            return

        token = auth_msg["token"]
        if token != SECRET_KEY:  # 使用你自己的 token 校验逻辑
            await websocket.send_json({"error": "Invalid token"})
            await websocket.close(code=4003)
            return

        # 记录连接建立
        logger.info(f"Client {client_id} is connecting")
        await manager.connect(client_id, websocket)

        while True:
            data = await websocket.receive_json()
            logger.info(f"Received message from {client_id}: {data}")
            message = {
                "client_id": client_id,
                "data": data
            }
            await manager.broadcast(json.dumps(message))
            await manager.handle_reply(client_id, data, db)
    except WebSocketDisconnect as e:
        logger.warning(f"Client {client_id} disconnected with code: {e.code}")
    except Exception as e:
        logger.error(f"Unexpected error for client {client_id}: {e}", exc_info=True)
    finally:
        manager.disconnect(client_id)
